# 修改后的代码（适用于 pandasai >= 3.x）
import duckdb
from pandasai import SmartDataframe
from pandasai.llm.openai import OpenAI


# 使用 DuckDB 读取和处理数据
# 示例：从 CSV 文件读取数据
def process_data_with_duckdb():
    # 连接到 DuckDB
    conn = duckdb.connect()

    # 从 CSV 文件加载数据到 DuckDB
    conn.execute("""
    CREATE TABLE titanic AS 
    SELECT * FROM 'test_cases/melb_data.csv'
    """)

    # 使用 SQL 查询处理数据
    result = conn.execute("""
    SELECT *
    FROM titanic
    """).df()

    # 关闭连接
    conn.close()

    return result


# 使用 PandasAI 进行自然语言分析
def analyze_with_pandasai(df):
    sdf = SmartDataframe(df, config={"llm": llm})  # 使用 SmartDataframe
    response = sdf.chat("用中文描述一下各个位置的房价分布，分析出哪些因素对房价影响较大，要求在300字内")
    return response


# 初始化 OpenAI LLM (请替换为你的 API 密钥)
llm = OpenAI(
    api_token="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")

# 处理数据
processed_data = process_data_with_duckdb()

# 分析数据
analysis_result = analyze_with_pandasai(processed_data)

print(analysis_result)

if __name__ == "__main__":
    print("DuckDB + PandasAI 集成成功！")
